import ssl
from unittest import result
import nltk
import string
import numpy as np
import os
from nltk.stem.porter import PorterStemmer

stemmer = PorterStemmer()


def pre_text(file):
    remove = str.maketrans('', '', string.punctuation)

    with open(file, errors='ignore') as f:
        lines = f.readlines()
        content = ''
        for line in lines:
            content += line.strip() + ''

    words_temp_1 = nltk.word_tokenize(content.lower().translate(remove).strip())

    words_temp_2 = [w for w in words_temp_1 if not (w in nltk.corpus.stopwords.words('english'))]

    words_temp_3 = [w for w in words_temp_2 if not w.isdigit()]

    content_cleaned = [stemmer.stem(w) for w in words_temp_3]

    return content_cleaned


def read_files(dir):
    categories = os.listdir(dir)

    # print(categories)

    contents_all = []

    for category in categories:

        category_articles = os.listdir(os.path.join(dir, category))

        for category_article in category_articles:
            content_cleaned = pre_text(os.path.join(dir, category, category_article))
            content_cleaned.append(category)
            # 将新闻类型添加至每个子类的最后一项
            contents_all.append(content_cleaned)

    # print(contents_all)

    return contents_all, categories


# 获取概率模型, 输入feat np.array格式 大小[N,D]
def trainPbmodel_X(feats):
    N, D = np.shape(feats)

    model = {}
    # 对每一维度的特征进行概率统计
    for d in range(D):
        data = feats[:, d].tolist()
        keys = set(data)
        N = len(data)
        model[d] = {}
        for key in keys:
            model[d][key] = float(data.count(key) / N)
    return model


# datas： list格式 每个元素表示1个特征序列
# labs：  list格式 每个元素表示一个文章类型
def trainPbmodel(datas, labs):
    # 定义模型
    model = {}
    # 获取分类的类别
    keys = set(labs)
    for key in keys:
        # 获得P(Y)
        Pbmodel_Y = labs.count(key) / len(labs)

        # 收集标签为Y的数据
        index = np.where(np.array(labs) == key)[0].tolist()

        feats = np.array(datas)[index]

        # 获得 P(X|Y)
        Pbmodel_X = trainPbmodel_X(feats)

        # 模型保存
        model[key] = {}
        model[key]["PY"] = Pbmodel_Y
        model[key]["PX"] = Pbmodel_X
    return model


# feat : list格式 一条输入特征
# model: 训练的概率模型
# keys ：考察标签的种类 
def getPbfromModel(feat, model, keys):
    results = {}
    eps = 0.00001
    for key in keys:

        py = model.get(key, eps).get("PY")

        # 分别获取 P(X|Y)
        model_X = model.get(key, eps).get("PX")
        list_px = []
        for d in range(len(feat)):
            pb = model_X.get(d, eps).get(feat[d], eps)
            list_px.append(pb)
        # 公式计算
        result = np.log(py) + np.sum(np.log(list_px))
        results[key] = result
    return results


# 将获取的样本预处理
# 将文章中的单词进行汇总
def reshapeData(dataSet):
    matrix = []
    for data in dataSet:
        matrix = matrix + data
    # 输出去重之前的长度
    print("特征数:" + str(len(matrix)))
    # 删除重复
    matrix = set(matrix)
    print("去除重复的特征数:" + str(len(matrix)))
    return matrix


if __name__ == '__main__':

    # 获取数据集
    dataSet, labels = read_files('./mini_newsgroups')

    print("获取训练集成功")

    matrix = reshapeData(dataSet)

    realM = list(matrix)

    data = []

    # 截取数据和标签
    datas = [i[:-1] for i in dataSet]
    labs = [i[-1] for i in dataSet]

    for d in datas:
        temp = []
        for i in matrix:
            temp.append(0)
        for i in range(len(d)):
            if (d[i] in realM):
                index = realM.index(d[i])
                temp[index] = 1
        data.append(temp)

    # 获取标签种类
    keys = set(labs)

    # 进行模型训练
    model = trainPbmodel(data, labs)

    print("模型建立成功")







    # 以下是测试代码

    # 获取测试新闻集
    testDataSets, testLables = read_files('./test')

    print("获取测试集成功")

    testDatas = [i[:-1] for i in testDataSets]
    testLabs = [i[-1] for i in testDataSets]

    # 所有测试集的特征集合
    testTempData = []

    # 逐个获取测试新闻的特征
    for d in testDatas:
        temp = []
        for i in matrix:
            temp.append(0)
        for i in range(len(d)):
            if d[i] in realM:
                index = realM.index(d[i])
                temp[index] = 1
        testTempData.append(temp)

    # 预测成功数
    successNum = 0

    for i in range(len(testTempData)):
        result = getPbfromModel(testTempData[i], model, keys)
        for key, value in result.items():
            if value == max(result.values()):
                if key == testLabs[i]:
                    successNum = successNum + 1
                    print(key + "预测成功")
                else:
                    print(key + "预测失败")

    print("预测成功率为 " + str(successNum / 20 * 10) + "%")

    # 用于训练集为20_newsgroups的测试结果输出
    # print("预测成功率为 "+str(successNum/200*100)+"%")
